package com.rain.rainspringboot3.encryption.strategy;

import com.rain.rainspringboot3.encryption.exception.DecryptionException;
import com.rain.rainspringboot3.encryption.exception.EncryptionException;
import jakarta.annotation.PostConstruct;
import org.apache.commons.lang3.StringUtils;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.Security;
import java.util.Base64;

import static com.rain.rainspringboot3.encryption.service.EncryptionServiceImpl.KEY_CACHE_MAP;

/**
 * 国密SM4加密策略实现
 * 使用SM4算法提供加密和解密功能
 */
@Component
public class SM4EncryptionStrategy implements EncryptionStrategy {

    static {
        // 添加BouncyCastle提供商
        Security.addProvider(new BouncyCastleProvider());
    }

    private static final String ALGORITHM = "SM4/CBC/PKCS5Padding";
    private static final int KEY_SIZE = 128;
    private static final int IV_LENGTH = 16;


    @Value("${encryption.strategies.sm4.key}")
    private String key;

    @Override
    public String encrypt(String content, String key) {
        try {
            if (StringUtils.isBlank(key)) {
                key = this.key;
            }
            // 生成随机IV
            byte[] iv = new byte[IV_LENGTH];
            java.security.SecureRandom random = new java.security.SecureRandom();
            random.nextBytes(iv);

            // 创建密钥规范
            SecretKeySpec keySpec = new SecretKeySpec(Base64.getDecoder().decode(key), "SM4");
            IvParameterSpec ivSpec = new IvParameterSpec(iv);

            // 初始化加密器
            Cipher cipher = Cipher.getInstance(ALGORITHM, "BC");
            cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);

            // 执行加密
            byte[] encrypted = cipher.doFinal(content.getBytes(StandardCharsets.UTF_8));

            // 组合IV和加密结果
            byte[] combined = new byte[iv.length + encrypted.length];
            System.arraycopy(iv, 0, combined, 0, iv.length);
            System.arraycopy(encrypted, 0, combined, iv.length, encrypted.length);

            return Base64.getEncoder().encodeToString(combined);
        } catch (Exception e) {
            throw new EncryptionException("SM4 encryption failed", e);
        }
    }

    @Override
    public String decrypt(String content, String key) {
        try {
            // 解码Base64
            byte[] combined = Base64.getDecoder().decode(content);

            // 分离IV和加密数据
            byte[] iv = new byte[IV_LENGTH];
            byte[] encrypted = new byte[combined.length - IV_LENGTH];
            System.arraycopy(combined, 0, iv, 0, iv.length);
            System.arraycopy(combined, iv.length, encrypted, 0, encrypted.length);

            // 创建密钥规范
            SecretKeySpec keySpec = new SecretKeySpec(Base64.getDecoder().decode(key), "SM4");
            IvParameterSpec ivSpec = new IvParameterSpec(iv);

            // 初始化解密器
            Cipher cipher = Cipher.getInstance(ALGORITHM, "BC");
            cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);

            // 执行解密
            byte[] decrypted = cipher.doFinal(encrypted);
            return new String(decrypted, StandardCharsets.UTF_8);
        } catch (Exception e) {
            throw new DecryptionException("SM4 decryption failed", e);
        }
    }

    @Override
    public String generateKey() {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance("SM4", "BC");
            keyGenerator.init(KEY_SIZE);
            SecretKey secretKey = keyGenerator.generateKey();
            return Base64.getEncoder().encodeToString(secretKey.getEncoded());
        } catch (Exception e) {
            throw new EncryptionException("Failed to generate SM4 key", e);
        }
    }

    @PostConstruct
    public void saveKey(){
        KEY_CACHE_MAP.put("SM4",key);
    }
}